import requests
import pandas as pd
from datetime import datetime
import io
from flask import Flask, request, jsonify, send_file

app = Flask(__name__, static_folder='.', static_url_path='')

@app.route('/')
def index():
    with open('index.html', 'r', encoding='utf-8') as f:
        return f.read()

@app.route('/get_data', methods=['POST'])
def get_data():
    data = request.json
    stock_code = data.get('stock_code', 'sh000001')  # 默认获取上证指数
    start_date = data.get('start_date')
    end_date = data.get('end_date')
    
    try:
        # 验证股票代码格式
        if not (stock_code.startswith('sh') or stock_code.startswith('sz')):
            return jsonify({'error': '股票代码格式错误，请使用sh或sz开头'}), 400
        
        # 转换股票代码格式为东方财富API所需格式
        market_code = '1' if stock_code.startswith('sh') else '0'
        eastmoney_code = stock_code[2:]  # 去掉sh或sz前缀
        secid = f"{market_code}.{eastmoney_code}"
        
        # 验证并格式化日期
        start_date_obj = datetime.strptime(start_date, '%Y-%m-%d')
        end_date_obj = datetime.strptime(end_date, '%Y-%m-%d')
        start_date_str = start_date_obj.strftime('%Y%m%d')
        end_date_str = end_date_obj.strftime('%Y%m%d')
        
        # 构建东方财富API请求URL
        url = f"http://push2his.eastmoney.com/api/qt/stock/kline/get?secid={secid}&fields1=f1,f2,f3,f4,f5&fields2=f51,f52,f53,f54,f55,f56,f57,f58&klt=101&fqt=1&beg={start_date_str}&end={end_date_str}"
        headers = {'User-Agent': 'Mozilla/5.0'}
        response = requests.get(url, headers=headers)
        data_json = response.json()
        
        if data_json['data'] is None:
            return jsonify({'error': '没有获取到数据，请检查股票代码和日期范围'}), 400
        
        klines = data_json['data']['klines']
        df_data = []
        for kline in klines:
            parts = kline.split(',')
            df_data.append({
                'Date': parts[0],
                'Open': float(parts[1]),
                'Close': float(parts[2]),
                'High': float(parts[3]),
                'Low': float(parts[4]),
                'Volume': float(parts[5])
            })
        
        df = pd.DataFrame(df_data)
        csv_buffer = io.StringIO()
        df.to_csv(csv_buffer, index=False, encoding='utf-8-sig')
        csv_buffer.seek(0)
        
        # 生成下载文件名，包含股票代码
        return send_file(
            io.BytesIO(csv_buffer.getvalue().encode('utf-8-sig')),
            mimetype='text/csv',
            as_attachment=True,
            download_name=f'{stock_code}_{start_date}_{end_date}.csv'
        )
    except ValueError as e:
        return jsonify({'error': f'日期格式错误: {str(e)}'}), 400
    except Exception as e:
        return jsonify({'error': f'数据获取失败: {str(e)}'}), 500

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5002)